/*

  Copyright (C) 2013 xubinbin 徐彬彬 (Beijing China)

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <http://www.gnu.org/licenses/>.

  
  filename: fa_ts2es.c
  version : v1.1.0
  time    : 2013/08/02 16:20 
  author  : xubinbin ( xubbwd@gmail.com )
  code URL: http://code.google.com/p/tsdemux/
  
*/

#include "fa_ps2es.h"

int psdemux_init(psdemux *psdemux_info)
{
	if(psdemux_info == NULL){
        printf("error psdemux_info\n");
        return -1;
	}
    
	psdemux_info->bufSize = PS_MAX_BUF;
    psdemux_info->buffer = malloc(psdemux_info->bufSize);
    memset(psdemux_info->buffer,0,psdemux_info->bufSize);
    psdemux_info->tmp_buffer = malloc(psdemux_info->bufSize);
    memset(psdemux_info->tmp_buffer,0,psdemux_info->bufSize);
    
	psdemux_info->cur_buf_Addr = 0;
	psdemux_info->find_ps_header_flag = 1;
	psdemux_info->pes_num = 0;
    psdemux_info->es_Data_Len = 0;
    psdemux_info->error_pes_packer = 0;
    psdemux_info->readin_data_flag = 0;
    
    return 0;
}


int find_ps_header(psdemux *psdemux_info)
{
    int stuffing_length = 0;
    int header_length = 0;
    psdemux_info->cur_buf_Addr = 0;
    //printf("file:%s line:%d\n",__FILE__,__LINE__);
    while (psdemux_info->cur_buf_Addr < psdemux_info->bufLen - 200)
    {
        //head of infomation processing;
        if ( psdemux_info->buffer[psdemux_info->cur_buf_Addr+0] == 0x00 &&
             psdemux_info->buffer[psdemux_info->cur_buf_Addr+1] == 0x00 &&
             psdemux_info->buffer[psdemux_info->cur_buf_Addr+2] == 0x01 &&
             psdemux_info->buffer[psdemux_info->cur_buf_Addr+3] == 0xBA) 
        {
            //printf("find ps start\n");
            //PS_pack_header_info
            psdemux_info->PS_pack_header_info.pack_stuffing_length = psdemux_info->buffer[psdemux_info->cur_buf_Addr+13] & 0x07;
            stuffing_length = psdemux_info->PS_pack_header_info.pack_stuffing_length;
            //printf("pack_stuffing_length = %d\n",psdemux_info->PS_pack_header_info.pack_stuffing_length);
            //printf("pack_stuffing_length = %d\n",stuffing_length);
            //PS_system_header_info
            //printf("psdemux_info->cur_buf_Addr = %d\n",psdemux_info->cur_buf_Addr);
            //printf("system_header_start_code[0] = 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+ 0+stuffing_length]);
            
            if ( psdemux_info->buffer[psdemux_info->cur_buf_Addr+14+stuffing_length] == 0x00 &&
                 psdemux_info->buffer[psdemux_info->cur_buf_Addr+15+stuffing_length] == 0x00 &&
                 psdemux_info->buffer[psdemux_info->cur_buf_Addr+16+stuffing_length] == 0x01 &&
                 psdemux_info->buffer[psdemux_info->cur_buf_Addr+17+stuffing_length] >> 4 == 0x0B ) 
            {
                psdemux_info->PS_system_header_info.header_length = 
                        ( psdemux_info->buffer[psdemux_info->cur_buf_Addr+18+stuffing_length] << 8 |
                          psdemux_info->buffer[psdemux_info->cur_buf_Addr+19+stuffing_length] ) & 0xffff;
                header_length = psdemux_info->PS_system_header_info.header_length;
                header_length += 6;
                //printf("header_length = %d\n",header_length);
            }
            psdemux_info->find_ps_header_flag = 1; 
            //printf("psdemux_info->cur_buf_Addr = %d\n",psdemux_info->cur_buf_Addr);
            //printf("psdemux_info->bufLen = %d\n",psdemux_info->bufLen);
            psdemux_info->cur_buf_Addr += (14+stuffing_length+header_length);
            psdemux_info->bufLen -=  psdemux_info->cur_buf_Addr;
            //printf("stuffing_length = %d\n",stuffing_length);
            //printf("header_length = %d\n",header_length);
            //printf("psdemux_info->cur_buf_Addr = %d\n",psdemux_info->cur_buf_Addr);
            //printf("psdemux_info->bufLen = %d\n",psdemux_info->bufLen);
            memcpy(psdemux_info->tmp_buffer,psdemux_info->buffer + psdemux_info->cur_buf_Addr,psdemux_info->bufLen);
            memset(psdemux_info->buffer,0,psdemux_info->bufSize);
            memcpy(psdemux_info->buffer,psdemux_info->tmp_buffer,psdemux_info->bufLen);
            memset(psdemux_info->tmp_buffer,0,psdemux_info->bufSize);
            //printf("file:%s line:%d\n",__FILE__,__LINE__);
            return 0;
        }
        else
        {
            psdemux_info->cur_buf_Addr++;
        }
    }
    return -1;
}

int process_pes_packet(psdemux *psdemux_info)
{
    int length_flag = 0;
    psdemux_info->cur_buf_Addr = 0;
    psdemux_info->es_Data_Len = 0;
    psdemux_info->error_pes_packer = 0;
    psdemux_info->PS_pes_info.PES_packet_length = 0;
    
    /*
    printf("buffer[length_flag +0] = 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+0]);
    printf("buffer[length_flag +1] = 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+1]);
    printf("buffer[length_flag +2] = 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+2]);
    printf("buffer[length_flag +3] = 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+3]);
    */
    
    while (psdemux_info->cur_buf_Addr < psdemux_info->bufLen - 200)
    {
        //head of infomation processing;
        if ( psdemux_info->buffer[psdemux_info->cur_buf_Addr+0] == 0x00 &&
             psdemux_info->buffer[psdemux_info->cur_buf_Addr+1] == 0x00 &&
             psdemux_info->buffer[psdemux_info->cur_buf_Addr+2] == 0x01 )
        {
            if(psdemux_info->buffer[psdemux_info->cur_buf_Addr+3] == 0xBA)
            {
                psdemux_info->find_ps_header_flag = 0; 
                break;
            }
            if(psdemux_info->buffer[psdemux_info->cur_buf_Addr+3] != 0xE0)  //E0 video  C0 audio
            {
                //printf("psdemux_info->buffer[psdemux_info->cur_buf_Addr+3] != 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+3]);
                psdemux_info->error_pes_packer = 1;
            }
            
            //printf("buffer[length_flag +0] = 0x%x\n",buffer[length_flag +0]);
            psdemux_info->PS_pes_info.packet_start_code_prefix = (psdemux_info->buffer[psdemux_info->cur_buf_Addr +0] << 16
                                                 | psdemux_info->buffer[psdemux_info->cur_buf_Addr +1] << 8
                                                 | psdemux_info->buffer[psdemux_info->cur_buf_Addr +2]) &0xffffff;
            psdemux_info->PS_pes_info.stream_id         = psdemux_info->buffer[psdemux_info->cur_buf_Addr +3] & 0xFF;
           
            //printf("buffer[length_flag +4] = 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+4]);
            //printf("buffer[length_flag +5] = 0x%x\n",psdemux_info->buffer[psdemux_info->cur_buf_Addr+5]);
            psdemux_info->PS_pes_info.PES_packet_length = ( (psdemux_info->buffer[psdemux_info->cur_buf_Addr +4] &0xFF) << 8 | 
                                                            (psdemux_info->buffer[psdemux_info->cur_buf_Addr +5]&0xFF) ) & 0xffff;
            if(psdemux_info->buffer[psdemux_info->cur_buf_Addr +6] >> 6 != 2 )
            {
                //printf("error file:%s,line:%d\n",__FILE__,__LINE__);
            }
            //printf("packet_pes.PES_packet_length = %d\n",psdemux_info->PS_pes_info.PES_packet_length);
            if((psdemux_info->PS_pes_info.PES_packet_length + psdemux_info->cur_buf_Addr + 6) > psdemux_info->bufLen)
                return 0;

            psdemux_info->PS_pes_info.PES_scrambling_control       = psdemux_info->buffer[psdemux_info->cur_buf_Addr +6] >> 4 & 0x03;
            psdemux_info->PS_pes_info.PES_priority                 = psdemux_info->buffer[psdemux_info->cur_buf_Addr +6] >> 3 & 0x01;
            psdemux_info->PS_pes_info.data_alignment_indicator     = psdemux_info->buffer[psdemux_info->cur_buf_Addr +6] >> 2 & 0x01;
            psdemux_info->PS_pes_info.copyright                    = psdemux_info->buffer[psdemux_info->cur_buf_Addr +6] >> 1 & 0x01;
            psdemux_info->PS_pes_info.original_or_copy             = psdemux_info->buffer[psdemux_info->cur_buf_Addr +6] >> 0 & 0x01;
            
            psdemux_info->PS_pes_info.PTS_DTS_flags                = psdemux_info->buffer[psdemux_info->cur_buf_Addr +7] >> 6 & 0x03;
            psdemux_info->PS_pes_info.ESCR_flag                    = psdemux_info->buffer[psdemux_info->cur_buf_Addr +7] >> 5 & 0x01;
            psdemux_info->PS_pes_info.ES_rate_flag                 = psdemux_info->buffer[psdemux_info->cur_buf_Addr +7] >> 4 & 0x01;
            psdemux_info->PS_pes_info.DSM_trick_mode_flag          = psdemux_info->buffer[psdemux_info->cur_buf_Addr +7] >> 3 & 0x01;
            psdemux_info->PS_pes_info.additional_copy_info_flag    = psdemux_info->buffer[psdemux_info->cur_buf_Addr +7] >> 2 & 0x01;
            psdemux_info->PS_pes_info.PES_CRC_flag                 = psdemux_info->buffer[psdemux_info->cur_buf_Addr +7] >> 1 & 0x01;
            psdemux_info->PS_pes_info.PES_extension_flag           = psdemux_info->buffer[psdemux_info->cur_buf_Addr +7] >> 0 & 0x01;
            psdemux_info->PS_pes_info.PES_header_data_length       = psdemux_info->buffer[psdemux_info->cur_buf_Addr +8];
            
            //printf("psdemux_info->PS_pes_info.PES_header_data_length = %d\n",psdemux_info->PS_pes_info.PES_header_data_length);
            length_flag += (9+psdemux_info->PS_pes_info.PES_header_data_length);
            //length_flag +=9;
            if(psdemux_info->PS_pes_info.PTS_DTS_flags == 0x02)
            {
                //length_flag += 5;//PTS
            }
            else if(psdemux_info->PS_pes_info.PTS_DTS_flags == 0x03)
            {
                //length_flag += 10;//PTS DTS
            }
            
            //printf("psdemux_info->pes_num = %d\n",psdemux_info->pes_num);
            //printf("pes_length = %d\n",psdemux_info->PS_pes_info.PES_packet_length);
            //printf("psdemux_info->cur_buf_Addr = %d\n",psdemux_info->cur_buf_Addr);
            psdemux_info->cur_buf_Addr += length_flag;
            psdemux_info->es_Data_Len = psdemux_info->PS_pes_info.PES_packet_length + 6 - length_flag;
            //printf("psdemux_info->es_Data_Len = %d\n",psdemux_info->es_Data_Len);
            psdemux_info->pes_num++;

            return 1;
        }
        else
        {
            //printf("@");
            //fflush(stdout);
            psdemux_info->cur_buf_Addr++;
        }
    }
    return 0;
}

int psdemux_get_es_process(psdemux *psdemux_info,unsigned char * input_buffer,int input_length,unsigned char * es_data,unsigned int *es_data_length)
{
    if(psdemux_info->readin_data_flag == 0)
    {
        if( (psdemux_info->bufLen + input_length) > psdemux_info->bufSize)
        {
            //reset
            psdemux_info->bufLen = 0;
            memset(psdemux_info->buffer,0,psdemux_info->bufSize);
            memcpy(psdemux_info->buffer,input_buffer,input_length);
            psdemux_info->bufLen = input_length;
            printf("psdemux_get_es_process buffer is full!\n");
            //data full; clear it; restart process
        }
        else
        {
            memcpy(psdemux_info->buffer + psdemux_info->bufLen,input_buffer,input_length);
            psdemux_info->bufLen += input_length;
        }
    }

    if(psdemux_info->find_ps_header_flag == 0)
    {
        if(find_ps_header(psdemux_info) < 0)
            return 0;
    }
    
    //pes parsing
    if(process_pes_packet(psdemux_info) > 0)
    {
        //output es data
        *es_data_length = psdemux_info->es_Data_Len;
        if(psdemux_info->es_Data_Len < 0)
            return 0;
        psdemux_info->bufLen = psdemux_info->bufLen - psdemux_info->cur_buf_Addr -(*es_data_length);
        memcpy(es_data, psdemux_info->buffer + psdemux_info->cur_buf_Addr, *es_data_length);
        
        memcpy(psdemux_info->tmp_buffer, 
               psdemux_info->buffer + psdemux_info->cur_buf_Addr +(*es_data_length),
               psdemux_info->bufLen);
        memset(psdemux_info->buffer,0,psdemux_info->bufSize);
        memcpy(psdemux_info->buffer,psdemux_info->tmp_buffer,psdemux_info->bufLen);
        memset(psdemux_info->tmp_buffer,0,psdemux_info->bufSize);
               
        psdemux_info->cur_buf_Addr = 0;  
        if(psdemux_info->error_pes_packer == 1)
        {
            //printf("error pes packer\n");
            return 0;
        }
        psdemux_info->readin_data_flag = 1; 
        return 1;
    }
    return 0;
}


int psdemux_deinit(psdemux *psdemux_info)
{
    if(psdemux_info->buffer != NULL)
    {
        free(psdemux_info->buffer);
        psdemux_info->buffer = NULL;
    }
    if(psdemux_info->tmp_buffer != NULL)
    {
        free(psdemux_info->tmp_buffer);
        psdemux_info->tmp_buffer = NULL;
    }
    if(psdemux_info != NULL)
    {
        free(psdemux_info);
        psdemux_info = NULL;
    }
    return 0;
}
